package tech.yozo.bamboo.model;

import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Paint;
import android.util.Log;

import com.baidu.paddle.fastdeploy.LitePowerMode;
import com.baidu.paddle.fastdeploy.RuntimeOption;
import com.baidu.paddle.fastdeploy.pipeline.PPOCRv3;
import com.baidu.paddle.fastdeploy.vision.OCRResult;
import com.baidu.paddle.fastdeploy.vision.ocr.Classifier;
import com.baidu.paddle.fastdeploy.vision.ocr.DBDetector;
import com.baidu.paddle.fastdeploy.vision.ocr.Recognizer;

import tech.yozo.bamboo.utils.Utils;

public class OcrEngine {
    private PPOCRv3 ocr;

    public void init(Context context) {
        if (ocr != null && ocr.initialized()) {
            return;
        }

        ocr = new PPOCRv3();

        int cpuThreadNum = 2;
        boolean enableLiteFp16 = false;
        LitePowerMode cpuPowerMode = LitePowerMode.LITE_POWER_HIGH;

        String realModelDir = context.getFilesDir() + "/models";
        String detModelName = "ch_PP-OCRv3_det_infer";
        String clsModelName = "ch_ppocr_mobile_v2.0_cls_infer";
        String recModelName = "ch_PP-OCRv3_rec_infer";

        String realDetModelDir = realModelDir + "/" + detModelName;
        String realClsModelDir = realModelDir + "/" + clsModelName;
        String realRecModelDir = realModelDir + "/" + recModelName;

        String srcDetModelDir = "models" + "/" + detModelName;
        String srcClsModelDir = "models" + "/" + clsModelName;
        String srcRecModelDir = "models" + "/" + recModelName;

        Utils.copyDirectoryFromAssets(context, srcDetModelDir, realDetModelDir);
        Utils.copyDirectoryFromAssets(context, srcClsModelDir, realClsModelDir);
        Utils.copyDirectoryFromAssets(context, srcRecModelDir, realRecModelDir);

        String realLabelPath = context.getFilesDir() + "/labels";
        Utils.copyDirectoryFromAssets(context, "labels", realLabelPath);

        String detModelFile = realDetModelDir + "/" + "inference.pdmodel";
        String detParamsFile = realDetModelDir + "/" + "inference.pdiparams";
        String clsModelFile = realClsModelDir + "/" + "inference.pdmodel";
        String clsParamsFile = realClsModelDir + "/" + "inference.pdiparams";
        String recModelFile = realRecModelDir + "/" + "inference.pdmodel";
        String recParamsFile = realRecModelDir + "/" + "inference.pdiparams";
        String recLabelFilePath = realLabelPath + "/" + "ppocr_keys_v1.txt";

        RuntimeOption detOption = new RuntimeOption();
        RuntimeOption clsOption = new RuntimeOption();
        RuntimeOption recOption = new RuntimeOption();
        detOption.setCpuThreadNum(cpuThreadNum);
        clsOption.setCpuThreadNum(cpuThreadNum);
        recOption.setCpuThreadNum(cpuThreadNum);
        detOption.setLitePowerMode(cpuPowerMode);
        clsOption.setLitePowerMode(cpuPowerMode);
        recOption.setLitePowerMode(cpuPowerMode);
        if (enableLiteFp16) {
            detOption.enableLiteFp16();
            clsOption.enableLiteFp16();
            recOption.enableLiteFp16();
        }

        DBDetector detModel = new DBDetector(detModelFile, detParamsFile, detOption);
        Classifier clsModel = new Classifier(clsModelFile, clsParamsFile, clsOption);
        Recognizer recModel = new Recognizer(recModelFile, recParamsFile, recLabelFilePath, recOption);
        if (ocr.init(detModel, clsModel, recModel)) {
            Log.d("OCR", "OCR初始化成功");
        }
    }

    public OcrText doOCR(Bitmap image) {
        if (ocr.initialized()) {
            // OCR 识别
            OcrText text = new OcrText();
            OCRResult result = ocr.predict(image);
            text.mTexts = result.mText;
            text.mBoxes = result.mBoxes;
            text.mRecScores = result.mRecScores;
            text.mClsScores = result.mClsScores;
            text.mClsLabels = result.mClsLabels;

            // TODO 用于测试
            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
//            text.mTexts = new String[]{"The future belongs to those who believe in the beauty of their dreams."};
//            text.mBoxes = new int[][]{{100, 1000, 2600, 1000, 2600, 1300, 100, 1300}};
            // 划线标注识别的内容
            if (text.mBoxes.length > 0) {
                // 准备绘制
                Canvas canvas = new Canvas(image);
                Paint paint = new Paint();

                for (int[] box : text.mBoxes) {
                    paint.setColor(Color.RED);
                    paint.setStrokeWidth(8);
                    canvas.drawLine(box[6], box[7], box[4], box[5], paint);
                    paint.setColor(Color.GREEN);
                    paint.setAlpha(90);
                    canvas.drawRect(box[0], box[1], box[4], box[5], paint);
                }
            }
            return text;
        }
        return null;
    }

    public void release() {
        if (ocr != null) {
            ocr.release();
        }
    }
}
